package org.apache.flink.connector;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.CheckpointListener;
import org.apache.flink.connector.elasticsearch.sink.Elasticsearch8AsyncSink;
import org.apache.flink.connector.elasticsearch.sink.Elasticsearch8AsyncSinkBuilder;
import org.apache.flink.streaming.api.environment.LocalStreamEnvironment;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import co.elastic.clients.elasticsearch.core.bulk.IndexOperation;
import org.apache.http.HttpHost;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.junit.jupiter.Testcontainers;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static org.assertj.core.api.AssertionsForClassTypes.assertThat;

/** Integration tests for {@link Elasticsearch8AsyncSink}. */
@Testcontainers
public class Elasticsearch8AsyncSinkITCase extends ElasticsearchSinkBaseITCase {
    private static boolean failed;

    @BeforeEach
    void setUp() {
        this.client = getRestClient();
        failed = false;
    }

    @AfterEach
    void shutdown() throws IOException {
        if (client != null) {
            client.close();
        }
    }

    @Test
    public void testWriteToElasticsearch() throws Exception {
        String index = "test-write-to-elasticsearch";

        try (StreamExecutionEnvironment env =
                     StreamExecutionEnvironment.getExecutionEnvironment().setParallelism(1)) {

            env.setRestartStrategy(RestartStrategies.noRestart());

            final Elasticsearch8AsyncSink<DummyData> sink =
                    Elasticsearch8AsyncSinkBuilder.<DummyData>builder()
                            .setMaxBatchSize(5)
                            .setHosts(
                                    new HttpHost(
                                            ES_CONTAINER.getHost(),
                                            ES_CONTAINER.getFirstMappedPort()))
                            .setElementConverter(
                                    (element, ctx) ->{
                                        List list = new ArrayList();
                                        list.add(new IndexOperation.Builder<>()
                                                .index(index)
                                                .id(element.getId())
                                                .document(element)
                                                .build());
                                        return list;
                                    }
                            )
                            .build();

            env.fromElements("first", "second", "third", "fourth", "fifth")
                    .map(
                            (MapFunction<String, DummyData>)
                                    value -> new DummyData(value + "_v1_index", value))
                    .sinkTo(sink);

            env.execute();
        }

        assertIdsAreWritten(index, new String[] {"first_v1_index", "second_v1_index"});
    }

    @Test
    public void testRecovery() throws Exception {
        String index = "test-recovery";

        try (final StreamExecutionEnvironment env = new LocalStreamEnvironment()) {

            env.enableCheckpointing(100L);

            final Elasticsearch8AsyncSink<DummyData> sink =
                    Elasticsearch8AsyncSinkBuilder.<DummyData>builder()
                            .setMaxBatchSize(5)
                            .setHosts(
                                    new HttpHost(
                                            ES_CONTAINER.getHost(),
                                            ES_CONTAINER.getFirstMappedPort()))
                            .setElementConverter(
                                    (element, ctx) ->{

                                        List list = new ArrayList();
                                        new IndexOperation.Builder<>()
                                                .index(index)
                                                .id(element.getId())
                                                .document(element)
                                                .build();
                                        return list;
                                    }
                                            )
                            .build();

            env.fromElements("first", "second", "third", "fourth", "fifth")
                    .map(
                            (MapFunction<String, DummyData>)
                                    value -> new DummyData(value + "_v1_index", value))
                    .map(new BuggyMapper())
                    .sinkTo(sink);

            env.execute();
        }

        assertThat(failed).isEqualTo(true);
    }

    private static class BuggyMapper
            implements MapFunction<DummyData, DummyData>, CheckpointListener {
        private int emittedRecords = 0;

        @Override
        public DummyData map(DummyData dummyData) throws InterruptedException {
            Thread.sleep(50);
            emittedRecords++;
            return dummyData;
        }

        @Override
        public void notifyCheckpointComplete(long l) throws Exception {
            if (!failed || emittedRecords != 0) {
                failed = true;
                throw new Exception();
            }
        }
    }

    /** DummyData is a POJO to helping during integration tests. */
    public static class DummyData {
        private final String id;

        private final String name;

        public DummyData(String id, String name) {
            this.id = id;
            this.name = name;
        }

        public String getId() {
            return id;
        }

        public String getName() {
            return name;
        }
    }
}
